package com.slipper.ai.flow.support;

import com.slipper.ai.flow.INodeDealer;
import com.slipper.ai.flow.vo.*;
import com.slipper.ai.vector.KnowledgeVectorService;
import com.slipper.util.Tool;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Component("aiKnowledgeNodeDealer")
public class KnowledgeNodeDealer implements INodeDealer {
    @Resource
    private KnowledgeVectorService knowledgeVectorService;

    @Override
    public String getType() {
        return NodeVO.TYPE_KNOWLEDGE;
    }

    @Override
    public List<LineVO> doDeal(FlowVO flow, NodeVO node, NodeInstVO nodeInst) {
        if (Tool.isBlank(nodeInst.getInputValue())) throw new RuntimeException("用户问题/输入变量不能为空！");
        if (!(node instanceof KnowledgeNodeVO knowledgeNode)) throw new RuntimeException("节点类型不正确！");
        if (Tool.isBlank(knowledgeNode.getKnowledgeList()) || knowledgeNode.getKnowledgeList().isEmpty()) {
            throw new RuntimeException("知识库不能为空！");
        }

        List<Document> documentList = new ArrayList<>();
        knowledgeNode.getKnowledgeList()
                .forEach(knowledge ->
                        documentList.addAll(knowledgeVectorService.similaritySearch(
                            knowledge.getId(),
                            nodeInst.getInputValue())
                ));

        Map<String, Object> outputMap = new HashMap<>();
        outputMap.put("文档列表", documentList);
        outputMap.put("文档内容", documentList
                .stream()
                .map(Document::getText)
                .collect(Collectors.joining("\n\n")));
        nodeInst.setOutput(outputMap);

        return flow.getLines()
                .stream()
                .filter(line -> line.getSource().equals(node.getId()))
                .collect(Collectors.toList());
    }
}
